
from codesecurity.feature.property_graph import build_const2identifier_graph_js,build_identifier2identifier_graph_js,build_identifier_mirror_graph_js
from codesecurity.feature.objects import AstDirection, CommonFeatureSet,Ast,AstPath
from codesecurity.feature.ast_image import AstImageBuilder
from codesecurity.utils.random_walk import random_walk_graph
from codesecurity.tasks.malicious_code_detect.meta import SuperParameter_MCD
from codesecurity.feature.node_type_const import JSNodeTypes

import random
import numpy as np
import string

class MCDFeatureBuilder:
    def __init__(self,sp:SuperParameter_MCD) -> None:
        self.sp=sp
        self.legal_chars_mapping={c:i for i,c in enumerate(self.sp.legal_chars)}
        
    
    def identifier_mapping(self,ast:Ast):
        identifier_name_mapping={}
        all_identifies=ast.search_nodes(['identifier'])
        
        for identifier in all_identifies:
            if identifier.ast_value not in identifier_name_mapping:
                identifier_name_mapping[identifier.ast_value]=f'v_{len(identifier_name_mapping)}{len(identifier_name_mapping)+1}'
        
        return identifier_name_mapping
        
    def build(self,common_feature:CommonFeatureSet,identifier_name_mapping=None):
        image_builder=AstImageBuilder(common_feature.ast_object)
        if identifier_name_mapping is None:
            identifier_name_mapping={}
        
        nodes,edges=build_const2identifier_graph_js(image_builder)
        _,i2i_edges=build_identifier2identifier_graph_js(image_builder)
        _,m2i_edges=build_identifier_mirror_graph_js(image_builder)
        
        edges=[e for e in edges if e[0] is not None and e[1] is not None]
        i2i_edges=[e for e in i2i_edges if e[0] is not None and e[1] is not None]
        m2i_edges=[e for e in m2i_edges if e[0] is not None and e[1] is not None]
        
        node_indexs=list(range(len(image_builder.ast_object.nodes)))
        edges_indexs=[(image_builder.node2index[edge[0]],image_builder.node2index[edge[1]]) for edge in edges]
        edges_indexs+=[(image_builder.node2index[edge[0]],image_builder.node2index[edge[1]]) for edge in i2i_edges]
        edges_indexs+=[(image_builder.node2index[edge[0]],image_builder.node2index[edge[1]]) for edge in m2i_edges]
        
        image=random_walk_graph(node_indexs,edges_indexs,self.sp.path_number,self.sp.step_number,self.sp.channel)
        dfg_images=np.zeros((self.sp.path_number,self.sp.channel,self.sp.step_number,self.sp.ast_node_dim),dtype='float32')

        for i in range(self.sp.path_number):
            for j in range(self.sp.step_number):
                for k in range(self.sp.channel):
                    node_index=int(image[i,j,k])
                    if node_index!=-1:
                        dfg_images[i,k,j,:]=self.node_vectorize(True,AstDirection.Down,image_builder.ast_object.nodes[node_index].ast_value)

        #image=image.reshape((self.sp.path_number,self.sp.step_number*self.sp.ast_node_dim,self.sp.channel))
        dfg_images=dfg_images.astype('float32')
        #image/=255
        
        
        return dfg_images
    
    def build_jsreaver(self,common_feature:CommonFeatureSet):
        image_builder=AstImageBuilder(common_feature.ast_object)
        identifier_name_mapping={}
        
        # path_number(512),node_number(256),[node_type,node_name]
        ret=[]
        
        identifier_by_name={}
        all_identifies=image_builder.search_nodes(['identifier'])
        for identifier in all_identifies:
            if identifier.ast_value not in identifier_by_name:
                identifier_by_name[identifier.ast_value]=[]
            identifier_by_name[identifier.ast_value].append(identifier)
        names=list(identifier_by_name.keys())
        
        
        if len(all_identifies)<2:
            return np.zeros((self.sp.path_number,self.sp.ast_path_max_length,self.sp.ast_node_dim),dtype='float32')
        
        for identifier in all_identifies:
            if identifier.ast_value not in identifier_name_mapping:
                identifier_name_mapping[identifier.ast_value]=f'v_{len(identifier_name_mapping)}{len(identifier_name_mapping)+1}'
        
        start2end_pairs=[]
        for _ in range(self.sp.path_number):
            name=random.choice(names)
            if len(identifier_by_name[name])<2:
                start2end_pairs.append(random.sample(all_identifies,2))
            else:
                start2end_pairs.append(random.sample(identifier_by_name[name],2))
        
        #start2end_pairs.sort(key=lambda x:image_builder.node2index[x[0]])
        
        for i,pair in enumerate(start2end_pairs):
            start,end=pair
            ast_path=image_builder.walk(start,end)
            
            ret.append(self.path_vectorize(ast_path,identifier_name_mapping))
    
    
        ret=np.stack(ret,axis=0)
        
        return ret

    def build_ast_image(self,common_feature:CommonFeatureSet,identifier_name_mapping=None):
        image_builder=AstImageBuilder(common_feature.ast_object)
        if identifier_name_mapping is None:
            identifier_name_mapping={}
        
        # path_number(512),node_number(256),[node_type,node_name]
        ret=np.zeros((self.sp.path_number,self.sp.ast_path_max_length,self.sp.ast_node_dim),dtype='float32')
        
        all_identifies=image_builder.search_nodes(['identifier'])
        
        
        if len(all_identifies)<2:
            all_identifies+=image_builder.search_nodes(['string'])
        
        if len(all_identifies)<2:
            return ret
        
        start2end_pairs=[]
        while len(start2end_pairs)<self.sp.path_number:
            start,end=random.sample(all_identifies,2)
            
            if image_builder.node2index[start]>image_builder.node2index[end]:
                start,end=end,start
            
            start2end_pairs.append((start,end))
        
        #start2end_pairs=list(start2end_pairs)
        start2end_pairs.sort(key=lambda x:image_builder.node2index[x[0]])
        
        for i,pair in enumerate(start2end_pairs):
            start,end=pair
            ast_path=image_builder.walk(start,end)
            
            ret[i,:,:]=self.path_vectorize(ast_path,identifier_name_mapping)
    
    
        
        return ret
    
    def path_vectorize(self,ast_path:AstPath,value_mapping:dict):
        ret=np.zeros((self.sp.ast_path_max_length,self.sp.ast_node_dim))
        
        i=0
        for is_leave,direction,value in ast_path.extract():
            value=value_mapping.get(value,value)
            ret[i,:]=self.node_vectorize(is_leave,direction,value)
            i+=1
            if i>= self.sp.ast_path_max_length:
                print(f'Warning: ast path length: {len(ret)} is too long. Expect {self.sp.ast_path_max_length} instead. ')
                break
        
        # if len(ret)>self.sp.ast_path_max_length:
        #     
        #     ret=ret[:self.sp.ast_path_max_length,:]
            
        return ret
    
    def node_vectorize(self,is_leave:bool,direction:int,value:str):
        #is_leave_mapping={True:0,False:1}
        #direction_mapping={AstDirection.start:0,AstDirection.Up:1,AstDirection.Down:2}
        
        ret=np.zeros(self.sp.ast_node_dim)
        
        
        # is_leave block length=2
        if is_leave:
            ret[0]=1
        else:
            ret[1]=1
            
        # direction block length=3
        ret[direction+2]=1
        
        for c in value:
            dev=self.legal_chars_mapping.get(c,len(self.legal_chars_mapping)-1)
            ret[dev+5]+=1
            
        return ret
        
class MCD_Feature:
    def __init__(self,image,label,origin=None) -> None:
        self.image=image
        self.label=label
        self.origin=origin

class MCD_BigFeatures:
    def __init__(self,features:list,label_mapping) -> None:
        self.features=features
        self.buffer=[]
        self.now_group_id=-1

        self.label_mapping=label_mapping

    def __getitem__(self,index):
        
        obj=[]
        for e in self.features:
            obj.append(e[index])
        
        return *obj,self.label_mapping[index]
    
    def __len__(self):
        return len(self.features[0])

class MCD_Features:
    def __init__(self,features) -> None:
        self.features=features

        self.label_mapping={}
        for e in self.features:
            if e.label not in self.label_mapping:
                self.label_mapping[e.label]=len(self.label_mapping)


    def label_all(self,label,mapping_id=0):
        for e in self.features:
            e.label=label
            
        self.label_mapping={label:mapping_id}
    
    def label_with_origin(self,mapping_handle=None):
        if mapping_handle is None:
            mapping_handle=lambda x:x.origin.path

        for e in self.features:
            e.label=mapping_handle(e)
            
        self.label_mapping={e.label:e.label for i,e in enumerate(self.features)}

    def __getitem__(self,index):
        obj=self.features[index]
        
        return *obj.image,self.label_mapping[obj.label]
    
    def __len__(self):
        return len(self.features)